// Copyright Epic Games, Inc. All Rights Reserved.

#include "PointCloudWorldPartitionHelpers.h"
#include "PointCloud.h" // for logging 
#include "PointCloudSliceAndDiceManager.h"
#include "Engine/World.h"
#include "WorldPartition/WorldPartition.h"
#include "WorldPartition/WorldPartitionEditorHash.h"

#if WITH_EDITOR
#include "Editor.h"
#include "SourceControlHelpers.h"
#include "PackageSourceControlHelper.h"
#include "FileHelpers.h"
#include "ObjectTools.h"
#endif

namespace PointCloudWorldPartitionHelpers
{

bool CheckoutManagedActors(UWorld* World, const TArray<TSoftObjectPtr<AActor>>& ActorsToCheckout)
{
	check(World);

	if (ActorsToCheckout.Num() == 0)
	{
		return true;
	}

#if WITH_EDITOR
	if (UWorldPartition* WorldPartition = World->GetWorldPartition())
	{
		TArray<FString> PackagesToCheckout;

		for (const TSoftObjectPtr<AActor>& ManagedActor : ActorsToCheckout)
		{
			if (const FWorldPartitionActorDescInstance* ActorDescInstance = WorldPartition->GetActorDescInstanceByPath(ManagedActor.ToSoftObjectPath()))
			{
				PackagesToCheckout.Emplace(ActorDescInstance->GetActorPackage().ToString());
			}
		}

		if (PackagesToCheckout.Num() > 0)
		{
			FPackageSourceControlHelper PackageHelper;
			if (!PackageHelper.Checkout(PackagesToCheckout))
			{
				UE_LOG(PointCloudLog, Warning, TEXT("Unable to checkout the required files from SCC"));
				return false;
			}
			else
			{
				return true;
			}
		}
	}
#endif

	return true;
}

bool DeleteManagedActors(UWorld* World, const TArray<TSoftObjectPtr<AActor>>& ActorsToDelete)
{
	check(World);

	if (ActorsToDelete.Num() == 0)
	{
		return true;
	}

#if WITH_EDITOR
	// Remove potential references to to-be deleted objects from the global selection sets.
	if (GIsEditor)
	{
		GEditor->ResetAllSelectionSets();
	}

	UWorldPartition* WorldPartition = World ? World->GetWorldPartition() : nullptr;

	if (WorldPartition)
	{
		TArray<FString> PackagesToDeleteFromSCC;
		TSet<UPackage*> PackagesToCleanup;

		for (const TSoftObjectPtr<AActor>& ManagedActor : ActorsToDelete)
		{
			const FWorldPartitionActorDescInstance* ActorDescInstance = WorldPartition->GetActorDescInstanceByPath(ManagedActor.ToSoftObjectPath());

			// Some actors don't have actor desc, but really they shouldn't be generated by RuleProcessor
			// However, it is likely that some actors don't exist anymore, so a null here is not a critical condition
			if (!ActorDescInstance)
			{
				continue;
			}

			// If actor is loaded, just remove from world and keep track of package to cleanup
			// Note that the WP update will be done in the CleanupAfterSuccessfulDelete automatically
			if (ActorDescInstance->GetActor())
			{
				PackagesToCleanup.Emplace(ActorDescInstance->GetActor()->GetExternalPackage());
				World->DestroyActor(ActorDescInstance->GetActor());
			}
			// Otherwise, if we're in a world partition, we can still delete the object without loading it
			else
			{
				PackagesToDeleteFromSCC.Emplace(ActorDescInstance->GetActorPackage().ToString());
				WorldPartition->RemoveActor(ActorDescInstance->GetGuid());
			}
		}

		// Save currently loaded packages so they get deleted
		if (PackagesToCleanup.Num() > 0)
		{
			ObjectTools::CleanupAfterSuccessfulDelete(PackagesToCleanup.Array(), /*bPerformReferenceCheck=*/true);
		}

		// Delete outstanding unloaded packages
		if (PackagesToDeleteFromSCC.Num() > 0)
		{
			FPackageSourceControlHelper PackageHelper;
			if (!PackageHelper.Delete(PackagesToDeleteFromSCC))
			{
				UE_LOG(PointCloudLog, Warning, TEXT("Unable to delete all files from SCC; deleted actors will come back on map reload."));
				return false;
			}
		}
	}
	else
#endif
	{
		// Not in editor, really unlikely to happen but might be slow
		for (const TSoftObjectPtr<AActor>& ManagedActor : ActorsToDelete)
		{
			if (ManagedActor.Get())
			{
				World->DestroyActor(ManagedActor.Get());
			}
		}
	}

	return true;
}

bool RevertUnchangedManagedActors(UWorld* World, const TArray<TSoftObjectPtr<AActor>>& ActorsToRevertUnchanged)
{
	check(World);

	if (ActorsToRevertUnchanged.Num() == 0)
	{
		return true;
	}

#if WITH_EDITOR
	if (UWorldPartition* WorldPartition = World->GetWorldPartition())
	{
		TArray<FString> PackagesToRevertUnchanged;

		for (const TSoftObjectPtr<AActor>& ManagedActor : ActorsToRevertUnchanged)
		{
			if (const FWorldPartitionActorDescInstance* ActorDescInstance = WorldPartition->GetActorDescInstanceByPath(ManagedActor.ToSoftObjectPath()))
			{
				PackagesToRevertUnchanged.Emplace(ActorDescInstance->GetActorPackage().ToString());
			}
		}

		if (PackagesToRevertUnchanged.Num() > 0)
		{
			return USourceControlHelpers::RevertUnchangedFiles(PackagesToRevertUnchanged, /*bSilent=*/false);
		}
	}
#endif

	return true;
}

bool GatherLoadedActors(UWorld* World, const TArray<TSoftObjectPtr<AActor>>& ActorsToProcess, TArray<TSoftObjectPtr<AActor>>& OutLoadedActors)
{
	bool bFoundLoadedOrInvalid = false;

#if WITH_EDITOR
	UWorldPartition* WorldPartition = World ? World->GetWorldPartition() : nullptr;
	if (WorldPartition)
	{
		for (const TSoftObjectPtr<AActor>& Actor : ActorsToProcess)
		{
			const FWorldPartitionActorDescInstance* ActorDescInstance = WorldPartition->GetActorDescInstanceByPath(Actor.ToSoftObjectPath());

			if (!ActorDescInstance)
			{
				bFoundLoadedOrInvalid = true;
			}
			else if (ActorDescInstance->GetActor())
			{
				OutLoadedActors.Add(Actor);
				bFoundLoadedOrInvalid = true;
			}
		}
	}
	else
#endif // WITH_EDITOR
	{
		OutLoadedActors.Append(ActorsToProcess);
		bFoundLoadedOrInvalid = true;
	}

	return bFoundLoadedOrInvalid;
}

bool GatherUnloadedActors(UWorld* World, const TArray<TSoftObjectPtr<AActor>>& ActorsToProcess, TArray<TSoftObjectPtr<AActor>>& OutUnloadedActors)
{
	bool bFoundUnloaded = false;

#if WITH_EDITOR
	UWorldPartition* WorldPartition = World ? World->GetWorldPartition() : nullptr;
	if (WorldPartition)
	{
		for (const TSoftObjectPtr<AActor>& Actor : ActorsToProcess)
		{
			const FWorldPartitionActorDescInstance* ActorDescInstance = WorldPartition->GetActorDescInstanceByPath(Actor.ToSoftObjectPath());
			if (ActorDescInstance && !ActorDescInstance->GetActor())
			{
				OutUnloadedActors.Add(Actor);
				bFoundUnloaded = true;
			}
		}
	}
#endif

	return bFoundUnloaded;
}

bool GetNewActorNameFromRecycledPackage(UWorld* World, TFunctionRef<TSoftObjectPtr<AActor>()> RecycleFn, FName& OutNewActorName)
{
	bool bRecycledPackage = false;

#if WITH_EDITOR
	UWorldPartition* WorldPartition = World ? World->GetWorldPartition() : nullptr;
	if (WorldPartition)
	{
		TSoftObjectPtr<AActor> ActorToReplace = RecycleFn();
		const FWorldPartitionActorDescInstance* ActorDescInstance = WorldPartition->GetActorDescInstanceByPath(ActorToReplace.ToSoftObjectPath());

		if (ActorDescInstance)
		{
			OutNewActorName = ActorDescInstance->GetActorName();
			// Very important: since we'll overwrite the file in WP, we must remove it first
			// and it will get added back on spawn
			WorldPartition->RemoveActor(ActorDescInstance->GetGuid());

			bRecycledPackage = true;
		}
	}
#endif

	return bRecycledPackage;
}

bool MoveNewActorsToRecycledPackages(UWorld* World, TArray<TSoftObjectPtr<AActor>>& NewActors, TFunctionRef<TSoftObjectPtr<AActor>()> RecycleFn, FBox& OutActorBounds)
{
	bool bMovedActor = false;

#if WITH_EDITOR
	UWorldPartition* WorldPartition = World ? World->GetWorldPartition() : nullptr;
	if (WorldPartition)
	{
		// Backup dirty state of level; it'll change due to the rename, but it's spurious
		bool bWorldWasDirtyBefore = World->GetOutermost()->IsDirty();

		// Try to associate any new actors to unclaimed actors
		TSet<UPackage*> PackagesToSave;
		TSet<UPackage*> PackagesToCleanup;

		for (TSoftObjectPtr<AActor>& ActorToRename : NewActors)
		{
			const FWorldPartitionActorDescInstance* ActorToRenameDescInstance = WorldPartition->GetActorDescInstanceByPath(ActorToRename.ToSoftObjectPath());

			TSoftObjectPtr<AActor> ActorToReplace = RecycleFn();
			const FWorldPartitionActorDescInstance* ActorToReplaceDescInstance = WorldPartition->GetActorDescInstanceByPath(ActorToReplace.ToSoftObjectPath());
			if (ActorToReplaceDescInstance && ActorToRenameDescInstance)
			{
				// Load actor so we can rename it to the new name/external package
				AActor* Actor = ActorToRename.LoadSynchronous();
				// Add the old package to the list of to be delete packages
				PackagesToCleanup.Emplace(Actor->GetExternalPackage());
				// Remove previous actor package from WP
				WorldPartition->RemoveActor(ActorToRenameDescInstance->GetGuid());
				// Remove old actor package from WP
				WorldPartition->RemoveActor(ActorToReplaceDescInstance->GetGuid());
				// Rename the actor, will add back the actor to WP
				Actor->Rename(*ActorToReplaceDescInstance->GetActorName().ToString(), nullptr);
				Actor->GetExternalPackage()->MarkAsFullyLoaded();
				PackagesToSave.Add(Actor->GetExternalPackage());

				// Update soft object path in data
				ActorToRename = Actor;

				OutActorBounds += Actor->GetComponentsBoundingBox(/*bNonColliding=*/true, /*bIncludeFromChildrenActors=*/true);
				bMovedActor = true;
			}
		}

		if (PackagesToSave.Num() > 0)
		{
			UEditorLoadingAndSavingUtils::SavePackages(PackagesToSave.Array(), true);
		}

		if (PackagesToCleanup.Num() > 0)
		{
			ObjectTools::CleanupAfterSuccessfulDelete(PackagesToCleanup.Array(), /*bPerformReferenceCheck=*/true);
		}

		// Put back original dirty state (if it wasn't before)
		if (!bWorldWasDirtyBefore)
		{
			World->GetOutermost()->SetDirtyFlag(false);
		}
	}
#endif // WITH_EDITOR

	return bMovedActor;
}

bool MoveNewActorsToRecycledPackages(UWorld* World, TArray<TSoftObjectPtr<AActor>>& NewActors, TFunctionRef<TSoftObjectPtr<AActor>()> RecycleFn)
{
	// Early out, nothing to do
	if (NewActors.Num() == 0)
	{
		return false;
	}

	FBox BoxToUnload(EForceInit::ForceInit);
	bool bMovedActors = MoveNewActorsToRecycledPackages(World, NewActors, RecycleFn, BoxToUnload);

	if (bMovedActors)
	{
#if WITH_EDITOR
		// Unload cells from WP based on the actor bounds
		if (World && World->GetWorldPartition())
		{
			UnloadRegion(World, BoxToUnload);
		}
#endif

		// Finally, garbage collect
		CollectGarbage(RF_NoFlags, true);
	}

	return bMovedActors;
}

void UnloadRegion(UWorld* World, const FBox& Box)
{
#if WITH_EDITOR
	if (World)
	{
		if (UWorldPartition* WorldPartition = World->GetWorldPartition())
		{
			if (WorldPartition->EditorHash)
			{
				WorldPartition->EditorHash->ForEachIntersectingActor(Box, [WorldPartition](FWorldPartitionActorDescInstance* ActorDescInstance)
				{
					FWorldPartitionHandle ActorHandle(WorldPartition, ActorDescInstance->GetGuid());
					FWorldPartitionHandlePinRefScope ActorPinRefScope(ActorHandle);
				});
			}
		}
	}
#endif
}

}